import numpy as np
import matplotlib.pyplot as plt
from scipy.ndimage import sobel

def create_synthetic_images():
    # 第一帧：5x5中心区域亮
    frame1 = np.array([
        [10, 10, 10, 10, 10],
        [10, 10, 200, 10, 10],
        [10, 10, 200, 10, 10],
        [10, 10, 200, 10, 10],
        [10, 10, 10, 10, 10]
    ], dtype=np.float32)
    
    # 第二帧：中心亮区向右上移动1像素
    frame2 = np.array([
        [10, 10, 10, 10, 10],
        [10, 10, 10, 10, 10],
        [10, 10, 10, 200, 10],
        [10, 10, 10, 200, 10],
        [10, 10, 10, 200, 10]
    ], dtype=np.float32)
    
    return frame1, frame2

def compute_gradients(frame1, frame2):
    Ix = sobel(frame1, axis=1, mode='constant')  # x方向梯度 [[9]]
    Iy = sobel(frame1, axis=0, mode='constant')  # y方向梯度
    It = frame2 - frame1                         # 时间梯度 [[3]]
    return Ix, Iy, It

def lucas_kanade(Ix, Iy, It, window_size=3):
    height, width = Ix.shape
    u = np.zeros((height, width))
    v = np.zeros((height, width))
    half_win = window_size // 2
    
    for i in range(half_win, height - half_win):
        for j in range(half_win, width - half_win):
            # 提取邻域窗口 [[7]]
            Ix_win = Ix[i-half_win:i+half_win+1, j-half_win:j+half_win+1].flatten()
            Iy_win = Iy[i-half_win:i+half_win+1, j-half_win:j+half_win+1].flatten()
            It_win = It[i-half_win:i+half_win+1, j-half_win:j+half_win+1].flatten()
            
            # 构建方程组 [[1]]
            A = np.column_stack((Ix_win, Iy_win))
            b = -It_win
            
            # 最小二乘解 [[4]]
            try:
                flow = np.linalg.lstsq(A, b, rcond=None)[0]
                u[i,j], v[i,j] = flow[0], flow[1]
            except np.linalg.LinAlgError:
                continue
                
    return u, v

# 可视化改进版
def visualize_flow(u, v):
    plt.figure(figsize=(8,6))
    plt.quiver(np.arange(u.shape[1]), np.arange(u.shape[0]), 
               u, -v,  # y轴方向反转以匹配图像坐标系
               angles='xy', scale_units='xy', scale=1, 
               color='r', width=0.0022)
    plt.title("Optical Flow Vectors (Right-Up Movement)")
    plt.xlabel("X axis")
    plt.ylabel("Y axis")
    plt.grid(True)
    plt.show()

# 主程序
if __name__ == "__main__":
    frame1, frame2 = create_synthetic_images()
    Ix, Iy, It = compute_gradients(frame1, frame2)
    u, v = lucas_kanade(Ix, Iy, It, window_size=3)
    
    print("Horizontal flow (u) at center (2,2):", u[2,2])
    print("Vertical flow (v) at center (2,2):", v[2,2])
    
    visualize_flow(u, v)